Tree diameter [ DFS, BFS]

Time: O(V + E); Space: O(E); medium

Given an undirected tree, return its diameter: the number of edges in a longest path in that tree.

The tree is given as an array of edges where edges[i] = [u, v] is a bidirectional edge between nodes u and v. Each node has labels in the set {0, 1, …, edges.length}.

Example 1:

Input: edges = [[0,1],[0,2]]

Output: 2

Explanation:

  • A longest path of the tree is the path 1 - 0 - 2.

Example 2:

Input: edges = [[0,1],[1,2],[2,3],[1,4],[4,5]]

Output: 4

Explanation:

  • A longest path of the tree is the path 3 - 2 - 1 - 4 - 5.

Notes:

  • 0 <= edges.length < 10^4

  • edges[i][0] != edges[i][1]

  • 0 <= edges[i][j] <= edges.length

  • The given edges form an undirected tree.

Hints:

  1. If it is a tree, then its nodes n = len(edges) + 1.

  2. Build the tree graph first. Then start DFS from node 0.

  3. DFS returns the deepest depth. DFS state needs current node, tree graph, and parent.

  4. For current node, for all its neighbors, as long as it is not parent, get the depth from it, pick 2 largest.

  5. Use these 2 largest to update the diameter.

  6. And returns the largest depth + 1.

1. DFS

Intuition For all nodes in the diameter, only one node is allowed that both left edge and right edge are used. So we could use a DFS to get this diameter.

[1]:

import collections

class Solution1(object):
    def treeDiameter(self, edges):
        """
        :type edges: List[List[int]]
        :rtype: int
        """
        if not edges:
            return 0

        adj = collections.defaultdict(list)
        for p, q in edges:
            adj[p].append(q)
            adj[q].append(p)
        visited = set()
        root = edges[0][0]
        visited.add(root)
        return max(self._helper(root, adj, visited))

    def _helper(self, node, adj, visited):
        visited.add(node)
        m1, m2, closed = 0, 0, 0
        for child in adj[node]:
            if child in visited:
                continue
            m, m_closed = self._helper(child, adj, visited)
            m += 1
            _, m1, m2 = sorted([m1, m2, m])
            closed = max(closed, m_closed)
        closed = max(closed, m1 + m2)
        return max(m1, m2), closed

[2]:

s = Solution1()
edges = [[0,1],[0,2]]
assert s.treeDiameter(edges) == 2
edges = [[0,1],[1,2],[2,3],[1,4],[4,5]]
assert s.treeDiameter(edges) == 4

[3]:

import collections

class Solution2(object):
    """
    Time:  O(|V| + |E|)
    Space: O(|E|)
    """
    def treeDiameter(self, edges):
        """
        :type edges: List[List[int]]
        :rtype: int
        """
        graph, length = collections.defaultdict(set), 0
        for u, v in edges:
            graph[u].add(v)
            graph[v].add(u)
        curr_level = {(None, u) for u, neighbors in graph.items() if len(neighbors) == 1}

        while curr_level:
            curr_level = {(u, v) for prev, u in curr_level
                          for v in graph[u] if v != prev}
            length += 1
        return max(length-1, 0)

[4]:

s = Solution2()
edges = [[0,1],[0,2]]
assert s.treeDiameter(edges) == 2
edges = [[0,1],[1,2],[2,3],[1,4],[4,5]]
assert s.treeDiameter(edges) == 4